# -*- coding: UTF-8 -*-

from __future__ import absolute_import, division, print_function, unicode_literals

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])

# Import the backtrader platform
import backtrader as bt

import MySQLDataFeed as mdf
import SimpleMySqlClass as smc
import talib as talib
import numpy as np
import akshare as ak
import pandas as pd
import json

# 1，输入code
# 2，检查沪深300，中正500是否在内
# 3，检查 t_stock_data_daily_${code} 是否建表，创建表
# 4，拉取code日线数据，落库
# 5，跑双均线模型回测
# 6，记录开仓买卖记录到流水表


def log(out_put):
    print("%s", out_put)


def get_stock_code_name(conn, hz300, zh500):
    code_name_list = []
    if hz300:
        sql_300 = f"select code,name from t_cfg_hs_300 limit 301"
        cfg_300_obj = conn.get_all(sql_300)
        if cfg_300_obj:
            # code_name_list += cfg_300_obj
            for code_name_tup in cfg_300_obj:
                code_name_list.append(
                    {
                        "code": code_name_tup[0],
                        "name": code_name_tup[1],
                        "type": "hz300",
                    }
                )

    if zh500:
        sql_500 = f"select code,name from t_cfg_zz_500  limit 501"
        cfg_500_obj = conn.get_all(sql_500)
        if cfg_500_obj:
            for code_name_tup in cfg_500_obj:
                code_name_list.append(
                    {
                        "code": code_name_tup[0],
                        "name": code_name_tup[1],
                        "type": "zh500",
                    }
                )
    # for code_name in code_name_list:
    #     print("300-500 tup = ", code_name)
    return code_name_list


# 是否是沪深300，中证500
def get_cfg_info(conn, code):
    sql_300 = f"select * from t_cfg_hs_300 where code='{code}' limit 1"
    cfg_300_obj = conn.get_one(sql_300)
    print("cfg_300_obj", cfg_300_obj)
    if cfg_300_obj:
        print(f"code=%{code},是沪深300成分股")
        return cfg_300_obj

    sql_500 = f"select * from t_cfg_zz_500 where code='{code}' limit 1"
    cfg_500_obj = conn.get_one(sql_500)
    if cfg_500_obj:
        print(f"code=%{code},是中证500成分股")
        return cfg_500_obj

    return None


# 检查+初始化建立表
def check_and_create_table(conn, code):
    base_table_name = "t_trade_data"
    table_name = base_table_name + "_" + code
    if not check_table_name(conn, table_name):
        create_like_table(conn, base_table_name, table_name)
    return


# 检查表是否存在
def check_table_name(conn, table_name):
    all_tables = conn.check_table(table_name)
    for t in all_tables:
        cur_table_name = t[0]
        print("table = ", cur_table_name)
        if cur_table_name == table_name:
            return True
    return False


# 创建表
def create_like_table(conn, base_table_name, table_name):
    copy_sql = f"create table {table_name} like {base_table_name}"
    rs = conn.execute(copy_sql)
    return rs


# 2, save data to mysql
def save_data_to_mysql(conn, stock_data_df):
    rs_ct = 0
    for index, row in stock_data_df.iterrows():
        # 插入新记录
        data = {
            "id": 0,
            "date": row["date"],
            "date_time": row["date_time"],
            "code": row["code"],
            "open": row["open"],
            "high": row["high"],
            "low": row["low"],
            "close": row["close"],
            "volume": row["volume"],
            "amount": row["amount"],
            "create_time": datetime.datetime.now(),
            "update_time": datetime.datetime.now(),
        }
        table_name = "t_trade_data_" + row["code"]
        rs = conn.insertIgnore(table_name, data)
        rs_ct = rs_ct + rs

    print("------------------------rs_ct=", rs_ct)
    return


# 3, query data from mysql
def read_data_from_mysql(conn, table_name):
    rs = conn.count(table_name, None)
    print("query rs = ", rs)
    pass


# 获取数据库链接
def get_mysql_conn():
    conn = smc.SimpleMySqlClass("127.0.0.1", "stocks", "haojg", "haojg")
    return conn


# 获取最近15个月的数据
def get_start_date_yyyyMMdd():
    end_date = datetime.datetime.now()
    # 一年+3个月
    start_date = end_date - datetime.timedelta(days=365 + 92)
    start_date_param = start_date.strftime("%Y%m%d")
    return start_date_param


# 当前时间为结束日期
def get_end_date_yyyyMMdd():
    end_date = datetime.datetime.now()
    end_date_param = end_date.strftime("%Y%m%d")
    return end_date_param


# 获取日线数据
def get_daily_data(code):
    start_date_param = get_start_date_yyyyMMdd()
    end_date_param = get_end_date_yyyyMMdd()

    stock_qfq_data = ak.stock_zh_a_hist(
        symbol=code, adjust="qfq", start_date=start_date_param, end_date=end_date_param
    )
    # print(stock_qfq_data)
    stock_qfq_data = stock_qfq_data.iloc[:, 0:7]

    # 日期     开盘     收盘     最高     最低     成交量  成交额
    # 列名记得这样定义好
    stock_qfq_data.columns = [
        "date",
        "open",
        "close",
        "high",
        "low",
        "volume",
        "amount",
    ]
    stock_qfq_data["date_time"] = stock_qfq_data.loc[:, ["date"]]
    stock_qfq_data["code"] = code
    # print(stock_qfq_data)
    return stock_qfq_data


def get_data_and_read(conn, code):
    # 获取链接
    # conn = get_mysql_conn()
    cfg_obj = get_cfg_info(conn, code)
    print("检查是否成分股=", cfg_obj)
    table_name = "t_trade_data_" + code
    code_desc = cfg_obj[3] if cfg_obj else "-"
    print(f"开始获取{ code_desc} 数据, code = {code}, table_name={table_name}")
    check_and_create_table(conn, code)
    # 拉取远端数据
    stock_qfq_data = get_daily_data(code)
    # 落库
    save_data_to_mysql(conn, stock_qfq_data)

    # 读数
    read_data_from_mysql(conn, table_name)
    # 关闭连接
    # conn.close()


def backtrader_test_by_data(code, name, isShowPlot):
    # Create a cerebro entity
    cerebro = bt.Cerebro()

    # 5.加载观察者 Observers
    # cerebro.addobserver(OrderObserver)

    # Add a strategy
    # cerebro.addstrategy(MyMaMacdRsi)
    cerebro.addstrategy(MyMaMacdRsi, code=code, name=name)
    # add kdj 指标
    # cerebro.addindicator(KDJ)

    tableName = "t_trade_data_" + code
    sql = f"select `date`, `open`, `high`, `low`, `close`, `volume` from {tableName}"

    data = mdf.MySQLDataFeed(sql)
    print("data = ", data)
    # Add the Data Feed to Cerebro
    cerebro.adddata(data, name=name)

    # Set our desired cash start
    # 初始资金 100,000,000
    cerebro.broker.setcash(100000.0)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)

    # Add a FixedSize sizer according to the stake
    cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    # Set the commission

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name="pnl")  # 返回收益率时序数据
    # 添加分析指标
    # 返回年初至年末的年度收益率
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name="_AnnualReturn")
    # 计算最大回撤相关指标
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name="_DrawDown")
    # 计算年化收益：日度收益
    cerebro.addanalyzer(bt.analyzers.Returns, _name="_Returns", tann=252)
    # 计算年化夏普比率：日度收益
    cerebro.addanalyzer(
        bt.analyzers.SharpeRatio,
        _name="_SharpeRatio",
        timeframe=bt.TimeFrame.Days,
        annualize=True,
        riskfreerate=0,
    )  # 计算夏普比率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio_A, _name="_SharpeRatio_A")
    # 返回收益率时序
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name="_TimeReturn")

    # Print out the starting conditions
    print("Starting Portfolio Value: %.2f" % cerebro.broker.getvalue())
    cerebro.broker.set_coc(True)  # 设置以当日收盘价成交
    # Run over everything
    results = cerebro.run()
    strat = results[0]
    # Print out the final result
    print("Final Portfolio Value: %.2f" % cerebro.broker.getvalue())
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # print(daily_return)

    # # 从返回的 result 中提取回测结果
    # strat = result[0]
    # # 返回日度收益率序列
    # # 打印评价指标
    print("AnnualReturn= ", strat.analyzers._AnnualReturn.get_analysis())
    print("SharpeRatio= ", strat.analyzers._SharpeRatio.get_analysis())
    print("DrawDown= ", strat.analyzers._DrawDown.get_analysis())
    print("_Returns= ", strat.analyzers._Returns.get_analysis())

    # Plot the result
    if isShowPlot:
        results = cerebro.plot(
            style="candle", barup="red", bardown="green", volume=True
        )


class KDJ(bt.Indicator):
    # J 快线；K 中线；D 慢线
    lines = ("K", "D", "J")

    params = (
        ("period", 24),
        ("period_dfast", 3),
        ("period_dslow", 3),
    )

    plotlines = dict(
        J=dict(
            _fill_gt=("K", ("red", 0.50)),
            _fill_lt=("K", ("green", 0.50)),
        )
    )

    def __init__(self):
        # Add a KDJ indicator
        self.kd = bt.indicators.StochasticFull(
            self.data,
            period=self.p.period,
            period_dfast=self.p.period_dfast,
            period_dslow=self.p.period_dslow,
        )

        self.lines.K = self.kd.percD
        self.lines.D = self.kd.percDSlow
        self.lines.J = self.K * 3 - self.D * 2


class MyMaMacdRsi(bt.Strategy):
    params = (
        ("period_ema5", 5),
        ("period_ema10", 10),
        ("period_ema20", 20),
        ("period_ema60", 60),
        ("code", "000000"),
        ("name", "------"),
    )

    # 打印日志
    def log(self, txt, dt=None):
        dt = dt or self.data.datetime.date(0)
        print("%s, %s" % (dt, txt))

    def stop(self):
        print("stop 停止")
        self.end_cash = self.broker.getcash()
        self.add_trade_log(self.start_cash, self.end_cash)
        self.conn.close()

    def start(self):
        print("start 开始")
        self.start_cash = self.broker.getcash()
        self.conn = get_mysql_conn()

    def __init__(self):
        # 用于保存订单
        self.order = None
        # 订单价格
        self.buyprice = None
        # 订单佣金
        self.buycomm = None

        # 定义变量保存所有收盘价
        self.dataopen = self.data.open
        self.datahigh = self.data.high
        self.datalow = self.data.low
        self.dataclose = self.data.close
        self.datavolume = self.data.volume
        self.datacode = self.params.code
        self.dataname = self.params.name
        # self.datacode = code
        # self.dataname = name
        self.datadatetime = self.data.datetime
        # self.datatradedatetime = num2date(int(self.data.datetime))
        print("data = ", self.data)
        # print("self.data = ", json.dumps(self.data.tradedatetime))
        # 计算5日均线
        # self.ema5 = bt.talib.EMA(self.datavolume, timeperiod=self.params.period_ema5)
        # 计算10日均线
        self.ema10 = bt.talib.EMA(self.dataclose, timeperiod=self.params.period_ema10)
        # 计算20日均线
        self.ema20 = bt.talib.EMA(self.dataclose, timeperiod=self.params.period_ema20)
        # print("ema20 = ", self.ema20)
        # # 计算60日均线
        self.ema60 = bt.talib.EMA(self.dataclose, timeperiod=self.params.period_ema60)

        self.datamacd = bt.indicators.MACDHisto(self.datas[0])

        # 相对强弱指数
        self.datarsi = bt.indicators.RSI(self.datas[0])
        # print("data ris = ", self.datarsi)

        self.datakdj = KDJ(self.datas[0])

        # # 顺势指标
        self.datacci = bt.talib.CCI(self.datahigh, self.datalow, self.dataclose)

        # # 资金流量指标
        self.datamfi = bt.talib.MFI(
            self.datahigh, self.datalow, self.dataclose, self.datavolume
        )

        # # 上升动向值
        self.datamom = bt.talib.MOM(self.dataclose, timeperiod=10)

        # 定义全局变量
        # self.count = 0
        for data in self.datas:
            # 转为tabib要求的数据格式
            opens = np.array(data.open.array)
            highs = np.array(data.high.array)
            lows = np.array(data.low.array)
            closes = np.array(data.close.array)
            # print(opens)
            # 数据放入self中

            # 计算十字星数据，结果为-100底部十字星，结果为100顶部十字星，0非十字星
            # res = talib.CDLDOJISTAR(opens, highs, lows, closes)
            # print("十字星数据", res)
            # self.star[data._id] = res

    def save_trade_log(self, row):
        print("保存总交易日志", row)
        # type, code, name, price, volume, remark, isNew
        data = {
            "id": 0,
            "type": row["type"],
            "code": row["code"],
            "name": row["name"],
            "price": row["price"],
            "position_price": row["position_price"],
            "volume": row["volume"],
            "amount": row["amount"],
            "trade_date": row["trade_date"],
            "remark": json.dumps(row["remark"]),
            "create_time": datetime.datetime.now(),
            "order_pnl": row["order_pnl"],
            "order_comm": row["order_comm"],
            "order_pnl_comm": row["order_pnl_comm"],
            "order_pnl_percentage": row["order_pnl_percentage"],
            "order_ref": row["order_ref"],
            "trade_ref": "",
        }
        table_name = "t_trade_log"
        rs = self.conn.insertIgnore(table_name, data)
        # print("insert trade log", rs)
        return rs

    def check_trade_position(self, code):
        print("check trade position")
        check_position_sql = (
            f"select volume from t_trade_position where code='{code}' limit 1"
        )
        check_pos_obj = self.conn.get_one(check_position_sql)
        print(" position = ", check_pos_obj)
        # 检查是否有仓位
        return check_pos_obj and check_pos_obj[0] and check_pos_obj[0] > 0

    def get_trade_position(self, code, trade_date):
        print("check trade position")
        check_position_sql = (
            f"select volume from t_trade_position where code='{code}' limit 1"
        )
        check_pos_obj = self.conn.get_one(check_position_sql)
        print(" position = ", check_pos_obj)
        if check_pos_obj and check_pos_obj[0] and check_pos_obj[0] > 0:
            return check_pos_obj[0]
        else:
            return 0

    def save_trade_position(self, row):
        print("save trade position")
        table_name = "t_trade_position"
        cur_code = self.datacode
        get_position_sql = f"select code,name,price,volume,amount from {table_name} where code='{cur_code}' limit 1"
        pos_obj = self.conn.get_one(get_position_sql)

        ori_volume = int(row["volume"])
        volume = row["volume"]
        amount = row["amount"]
        price = row["price"]
        trade_type = row["type"]

        # 卖出,负仓位
        if trade_type == "sell":
            volume = 0 - ori_volume

        if pos_obj:
            volume += int(pos_obj[3])

            update_condition = " code = '" + str(cur_code) + "'"
            update_data = {
                "volume": volume,
                "trade_date": row["trade_date"],
                "update_time": datetime.datetime.now(),
            }
            rs = self.conn.update(table_name, update_data, update_condition)
            log_msg = f"有仓位，更新type={trade_type}price={price:.2f},volume={ori_volume},amount={amount:.2f} ,执行结果:{rs}"
            self.log(log_msg)

        else:

            insert_data = {
                "id": 0,
                "code": row["code"],
                "name": row["name"],
                "volume": volume,
                "trade_date": row["trade_date"],
                "update_time": datetime.datetime.now(),
                "create_time": datetime.datetime.now(),
            }

            rs = self.conn.insert(table_name, insert_data)
            log_msg = f"没有仓位，新增type={trade_type}price={price:.2f},volume={ori_volume},amount={amount:.2f},执行结果:{rs}"
            self.log(log_msg)

        return

    def update_order_info(self, update_data, order_ref):
        table_name = "t_trade_log"
        cur_code = self.datacode
        find_one_sql = f"select id,order_ref,code from {table_name} where code='{cur_code}' and order_ref='{order_ref}' order by id desc limit 1"
        rs = self.conn.get_one(find_one_sql)
        if rs and rs[0]:
            condition = " id = " + str(rs[0])
            self.conn.update(table_name, update_data, condition)

        return

    def notify_order(self, order):
        # print("order = ", order)
        # print("order.executed = ", order.executed)
        # 等待订单提交、订单被cerebro接受
        if order.status in [order.Submitted, order.Accepted]:
            return

        # 等待订单完成
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    "买 EXECUTED, ref:%.0f，开仓价格: %.2f, 花费: %.2f, 手续费: %.2f, 数量: %.2f, 股票: %s"
                    % (
                        order.ref,  # 订单编号
                        order.executed.price,  # 成交价
                        order.executed.value,  # 成交额
                        order.executed.comm,  # 佣金
                        order.executed.size,  # 成交量
                        order.data._name,
                    )
                )
                update_data = {"order_comm": order.executed.comm}
                self.update_order_info(update_data, order.ref)

            else:
                self.log(
                    "卖 EXECUTED, ref:%.0f,卖出价格: %.2f, 收回资金: %.2f, 手续费: %.2f, 毛利润：%2f, 数量：%.2f，股票：%s"
                    % (
                        order.ref,
                        order.executed.price,
                        order.executed.value,
                        order.executed.comm,
                        order.executed.pnl,
                        order.executed.size,
                        order.data._name,
                    )
                )

                update_data = {
                    "order_pnl": order.executed.pnl,
                    "order_comm": order.executed.comm,
                    "order_pnl_comm": order.executed.pnl - order.executed.comm,
                }
                self.update_order_info(update_data, order.ref)
                self.order = None

        # 如果订单保证金不足，将不会完成，而是执行以下拒绝程序
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log("拒绝执行，Order Canceled/Margin/Rejected = ", order.status)

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(
            f"策略收益：编号:{trade.ref:.0f}, 毛收益 {trade.pnl:.2f}, 净收益 {trade.pnlcomm:.2f}"
        )

    # 购买，建立持仓
    def buy_position(self, can_use_cash):
        cur_can_buy_shou = int(can_use_cash / (self.dataclose[0] * 100))
        buy_size = cur_can_buy_shou * 100
        if buy_size <= 0:
            log_msg = f"2本金不够一手，放弃购买。当前可用本金： {can_use_cash:.2f}, 当前股价： {self.dataclose[0]:.2f}"
            self.log(log_msg)
            return
        self.order = self.buy(size=buy_size)
        self.add_buy_log(buy_size)
        return

    # 添加购买交易日志
    def add_buy_log(self, buy_size):
        buy_row = {
            "code": self.datacode,
            "name": self.dataname,
            "type": "buy",
            "volume": buy_size,
            "price": self.dataclose[0],
            "position_price": self.dataclose[0],
            "amount": self.dataclose[0] * buy_size,
            "trade_date": bt.num2date(self.datadatetime[0]),
            "order_pnl": 0,
            "order_comm": 0,
            "order_pnl_comm": self.order.executed.comm,
            "order_pnl_percentage": 0,
            "order_ref": self.order.ref,
            "remark": {
                "ema20": self.ema20[0],
                "pre_ema20": self.ema20[-1],
                "close": self.dataclose[0],
                "pre_close": self.dataclose[1],
                "rsi": self.datarsi[0],
                "pre_rsi": self.datarsi[-1],
                "macd_histo": self.datamacd.histo[0],
                "pre_macd_histo": self.datamacd.histo[-1],
                "macd": self.datamacd.macd[0],
                "pre_macd": self.datamacd.macd[-1],
            },
        }
        self.save_trade_log(buy_row)
        self.save_trade_position(buy_row)
        return

    def add_trade_log(self, start_cash, end_cash):
        trade_row = {
            "code": self.datacode,
            "name": self.dataname,
            "type": "trade",
            "price": end_cash,
            "position_price": start_cash,
            "amount": start_cash,
            "trade_date": bt.num2date(self.datadatetime[0]),
            "volume": 0,
            "order_pnl": 0,
            "order_comm": 0,
            "order_pnl_comm": 0,
            "order_pnl_percentage": 100 * (end_cash - start_cash) / start_cash,
            "order_ref": 0,
            "remark": {},
        }
        self.save_trade_log(trade_row)
        return

    # 卖出持仓
    def sell_position(self, sell_size):
        if sell_size == -1:
            sell_size = self.broker.getposition(self.data).size

        self.order = self.sell(size=sell_size)
        self.add_sell_log(sell_size)
        return

    # 添加卖出交易日志
    def add_sell_log(self, sell_size):

        if sell_size <= 0:
            print("没有仓位不需要卖出")
            return

        # 当前持仓量
        cur_position_size = self.broker.getposition(self.data).size
        # 当前持仓成本
        cur_position_price = self.broker.getposition(self.data).price
        cur_order_pnl = (self.dataclose[0] - cur_position_price) * cur_position_size
        cur_order_pnl_comm = cur_order_pnl - self.order.executed.comm
        cur_order_pnl_percentage = (
            100 * (self.dataclose[0] - cur_position_price) / (cur_position_price)
        )
        sell_row = {
            "code": self.datacode,
            "name": self.dataname,
            "type": "sell",
            "volume": sell_size,
            "price": self.dataclose[0],
            "position_price": cur_position_price,
            "amount": self.dataclose[0] * sell_size,
            "trade_date": bt.num2date(self.datadatetime[0]),
            "order_pnl": cur_order_pnl,
            "order_comm": self.order.executed.comm,
            "order_pnl_comm": cur_order_pnl_comm,
            "order_pnl_percentage": cur_order_pnl_percentage,
            "order_ref": self.order.ref,
            "remark": {
                "ema20": self.ema20[0],
                "pre_ema20": self.ema20[-1],
                "close": self.dataclose[0],
                "pre_close": self.dataclose[-1],
                "rsi": self.datarsi[0],
                "pre_rsi": self.datarsi[-1],
                "macd_histo": self.datamacd.histo[0],
                "pre_macd_histo": self.datamacd.histo[-1],
                "macd": self.datamacd.macd[0],
                "pre_macd": self.datamacd.macd[-1],
            },
        }
        self.save_trade_log(sell_row)
        self.save_trade_position(sell_row)
        return

    def get_position_from_db(self):
        return

    # 判断历史数据是不是结束了.结束则平仓所有持仓
    def is_bar_stop(self, data, last_num):
        # print("len(data) = ", len(data))
        # print("data.buflen() = ", data.buflen())
        if len(data) >= (data.buflen() - last_num):
            if self.getposition(data).size != 0:
                # self.close(data)
                print("------最后一根K线了,还有仓位 ------")
            return True
        else:
            return False

    # 策略逻辑实现
    def next(self):
        # 最后一根K线，清仓
        if self.is_bar_stop(self.data, 1):
            print("结束，最后2根K线了,如果有仓位，清仓")
            self.sell_position(-1)
            return

        # self.check_trade_position(self.datacode)
        # 当前可用资金
        cur_cash = self.broker.getcash()
        # 当前总资产
        # cur_position_value = self.broker.getvalue()
        # 当前持仓量
        cur_position_size = self.broker.getposition(self.data).size
        # 当前持仓成本
        # cur_position_price = self.broker.getposition(self.data).price
        # print("当前可用资金", cur_cash)
        # print("当前总资产", self.broker.getvalue())
        # print("当前持仓量", self.broker.getposition(self.data).size)
        # print("当前持仓成本", self.broker.getposition(self.data).price)
        # print("self.datakdj.k = ", self.datakdj.K[0])

        # 跌破10日均线
        cur_down_cross_10_line = self.dataclose[0] < self.ema10[0]
        cur_down_cross_20_line = self.dataclose[0] < self.ema20[0]
        cur_down_cross_60_line = self.dataclose[0] < self.ema60[0]

        # 10日均线小于20日
        cur_ema_10_lt_20 = self.ema10[0] < self.ema20[0]
        # 20日均线大于60日
        cur_ema_20_gt_60 = self.ema20[0] > self.ema60[0]
        # 20日线拐头向上
        cur_ema_20_gt_pre_20 = self.ema20[0] >= self.ema20[-1]
        # 60日线拐头向上
        cur_ema_60_gt_pre_60 = self.ema60[0] >= self.ema60[-1]

        # macd 金叉，macd零轴之上，macd上升趋势
        cur_macd_jincha = self.datamacd.histo[0] > 0
        cur_macd_gt_0 = self.datamacd.macd[0] >= 0
        cur_macd_up = self.datamacd.macd[0] > self.datamacd.macd[-1]

        # rsi 大于50，rsi上升趋势
        cur_rsi_gt_50 = self.datarsi[0] > 50
        cur_rsi_up = self.datarsi[0] >= self.datarsi[-1]
        # K线当日收阳线，大于昨日上升趋势，
        cur_k_line_up = self.dataclose > self.dataopen

        # 放量
        cur_volume_up = self.datavolume[0] > self.datavolume[-1]

        # 趋势上升
        cur_ci_up = self.datacci[0] > 0 and self.datacci[0] > self.datacci[-1]

        cur_k_qushi = (
            self.dataclose[0] > self.dataclose[-1]
            and self.datahigh[0] > self.datahigh[-1]
            and self.datalow[0] > self.datalow[-1]
        )
        cur_ema_qushi = cur_ema_20_gt_pre_20 and cur_ema_60_gt_pre_60

        cur_jiang = self.dataclose[0] < self.datalow[-1]

        # macd 金叉，上行；rsi强势
        is_kai_cang = (
            cur_macd_jincha
            and cur_macd_up
            and cur_rsi_gt_50
            and cur_k_line_up
            and cur_k_qushi
            and cur_volume_up
            and cur_ci_up
            and cur_ema_qushi
            and cur_ema_20_gt_60
            and cur_macd_gt_0
        )

        # 当今天的20日均线大于60日均线并且昨天的20日均线小于60日均线，则进入市场（买）
        # if self.ema20[0] > self.ema60[0] and self.ema20[-1] < self.ema60[-1]:
        # if (
        #     cur_ema_20_gt_60
        #     and cur_macd_jincha
        #     and cur_macd_gt_0
        #     and cur_macd_up
        #     and cur_rsi_gt_50
        #     and cur_k_line_up
        #     and cur_k_qushi
        # ):
        if is_kai_cang:
            self.log("买点：")
            # print("可以考虑买入，self.order = ", self.order)
            # 判断订单是否完成，完成则为None，否则为订单信息
            if not self.order is None:
                # self.log("已经存在订单了")
                return
            # cur_can_buy_shou = int(cur_cash / (self.dataclose[0] * 100))
            # buy_size = cur_can_buy_shou * 100
            # post_total = self.dataclose[0] * buy_size
            # # print("需要花费：", post_total, "元")
            # # 若上一个订单处理完成，可继续执行买入操作
            # self.order = self.buy(size=buy_size)
            self.buy_position(cur_cash)
            # print("已经买入，self.order = ", self.order)
            # print("已经买入 close_price=", self.dataclose[0], " ok!!!")

        # elif cur_down_cross_10_line:
        elif cur_down_cross_20_line:

            # print("需要卖出，self.order = ", self.order)
            if self.order is None:
                # self.log("没有订单需要卖出")
                return

            # 卖出(全部卖出)
            sell_size = cur_position_size
            self.sell_position(sell_size)


def test_back_trade():
    try:
        print("start test ...")
        conn = get_mysql_conn()
        code_name_type_list = get_stock_code_name(conn, hz300=True, zh500=True)
        for code_name in code_name_type_list:
            code = code_name["code"]
            name = code_name["name"]
            type = code_name["type"]
            print(f"code={code}, name={name}, type={type}")
            # print("code-name", code_name["code"])
            # get_data_and_read(conn, code)
            backtrader_test_by_data(code, name, False)
    finally:
        conn.close()
        print("end test.....")


def get_fund_stock_to_db(conn):
    """
    data =
     序号               100
     股票代码          600522
     股票简称            中天科技
     报告期       2024-03-31
     基金覆盖家数           126
     持股总数       308658355
     持股总市值      433047.67
     Name: 99, dtype: object
    """
    # date="20210630"; choice of {"XXXX0331", "XXXX0630", "XXXX0930", "XXXX1231"}, 其中 XXXX 为年份
    fund_report_stock_info_df = ak.fund_report_stock_cninfo(date="20240331")
    # print(fund_report_stock_cninfo_df)
    # need_fund = fund_report_stock_cninfo_df[:5]
    # for index, row in need_fund.iterrows():
    #     print("data = ", row[0], row[1], row[2], row[3], row[4])
    #     print(
    #         "data = ",
    #         row["序号"],
    #         row["股票代码"],
    #         row["股票简称"],
    #         row["报告期"],
    #         row["基金覆盖家数"],
    #     )

    rs_ct = 0
    for index, row in fund_report_stock_info_df.iterrows():
        # 插入新记录
        data = {
            "id": 0,
            "order_no": row["序号"],
            "code": row["股票代码"],
            "name": row["股票简称"],
            "report_date": row["报告期"],
            "fun_com_num": row["基金覆盖家数"],
            "position_volume": row["持股总数"],
            "position_value": row["持股总市值"],
            "create_time": datetime.datetime.now(),
        }
        table_name = "t_fund_stock_info"
        rs = conn.insertIgnore(table_name, data)
        rs_ct = rs_ct + rs

    print("---------------基金---------=", rs_ct)
    return


if __name__ == "__main__":
    # '600519', '贵州茅台'
    # 300438	鹏辉能源
    # 伊利股份[600887]
    # code = "000333"
    # name = "美的集团"
    code = "600887"
    name = "伊利股份"
    conn = get_mysql_conn()
    # # code = "002594"
    get_data_and_read(conn, code)
    conn.close()
    
    backtrader_test_by_data(code, name, True)
    # test_back_trade()

    # conn = get_mysql_conn()
    # get_fund_stock_to_db(conn)
    # conn.close()
